package com.atguigu.flink.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.OperatorStateStore;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

/**
 * Created by Smexy on 2023/4/10
 * 
 *
 *  UnionListState: 也是 operatorState
 *                  和ListState是亲戚，但是区别在于数据在恢复时，是平均分配还是每个Task获得全量。
 *
 *
 *       备份的数据是: [a,b,c,d,e,f]
 *
 *       使用 ListState 恢复，它是平均分配。
 *              恢复的Task的并行度是2，那么此时
 *                      Task1:  [a,c,e]
 *                      Task2:  [b,d,f]
 *
 *       使用 UnionListState 恢复，它是每个Task获取备份数据的全部。
 *                恢复的Task的并行度是2，那么此时
 *                        Task1:  [a,b,c,d,e,f]
 *                        Task2:  [a,b,c,d,e,f]
 *
 *  --------------------------
 *      KafkaSource使用的是这个状态。
 *
 *          消费者组: test1
 *              task1: 消费者线程
 *                      消费 topicA-0,200
 *                      消费 topicA-2,100
 *              task2: 消费者线程
 *                      消费 topicA-1,300
 *                      消费 topicA-3,400
 *
 *           source算子保存的状态就是offset。
 *
 *       程序挂掉了，重启了，改了并行度，重新创建KafkaSource，之后会恢复状态
 *              触发kafka的rebalce(重分配)，使用 UnionListState可以让每个Task获取到当前组所消费的offset的全部数据。
 *
 *              消费者组: test1
 *                  task1: 消费者线程  topicA-0,200
 *                                   topicA-3,400
 *
 *                  task2: 消费者线程  topicA-1,300
 *
 *                  task3: 消费者线程  topicA-2,100
 *
 *---------------------
 *  ListState:
 *                  add(e): 添加一个元素
 *                  addAll(List x): 添加一个集合
 *                  update(List x): 覆盖添加x集合
 *                  clear(): 清空集合
 *                  get(): 获取集合中的元素
 *
 *
 *
 *
 *
 */
public class Demo2_UnionState
{
    public static void main(String[] args) {
        
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(2);
        env.enableCheckpointing(2000);
        
                env
                   .socketTextStream("hadoop102", 8888)
                   .map(new MyMapFunction())
                   .addSink(new SinkFunction<String>()
                   {
                       @Override
                       public void invoke(String value, Context context) throws Exception {
                           if (value.contains("x")){
                               //模拟抛异常
                               throw new RuntimeException("出异常了...");
                           }
                           System.out.println(value);
                       }
                   });

        
                try {
                            env.execute();
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
        
    }

    public static class MyMapFunction implements MapFunction<String,String>, CheckpointedFunction
    {

        //使用flink提供的状态，就当一个List用
        private ListState<String> result;

        @Override
        public String map(String value) throws Exception {
            result.add(value);
            return result.get().toString();
        }


        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            System.out.println("MyMapFunction.snapshotState");
        }


        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            System.out.println("MyMapFunction.initializeState");

            OperatorStateStore operatorStateStore = context.getOperatorStateStore();

            ListStateDescriptor<String> listStateDescriptor = new ListStateDescriptor<>("strs", String.class);
            //result = operatorStateStore.getListState(listStateDescriptor);
            result = operatorStateStore.getUnionListState(listStateDescriptor);
        }
    }
}
